import os
import json
import base64
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from Crypto.Random import get_random_bytes

class KeyManager:
    """管理加密密钥和IV，包括生成、导入/导出和存储。"""
    
    def __init__(self, keys_file="keys.json"):
        """
        初始化密钥管理器。
        
        参数:
            keys_file (str): 密钥存储文件的路径
        """
        self.keys_file = keys_file
        self.keys = self._load_keys()
    
    def _load_keys(self):
        """从文件加载保存的密钥和IV。"""
        try:
            if os.path.exists(self.keys_file):
                with open(self.keys_file, 'r', encoding='utf-8') as f:
                    return json.load(f)
            return {}
        except Exception:
            return {}
    
    def _save_keys(self):
        """将密钥和IV保存到文件。"""
        with open(self.keys_file, 'w', encoding='utf-8') as f:
            json.dump(self.keys, f, indent=2, ensure_ascii=False)
    
    def generate_key(self, key_size=32, key_name=None):
        """
        生成新的随机密钥。
        
        参数:
            key_size (int): 密钥大小（字节）（AES为16、24或32；DES为8；3DES为24）
            key_name (str): 保存密钥的可选名称
            
        返回:
            bytes: 生成的密钥
        """
        key = get_random_bytes(key_size)
        if key_name:
            self.save_key(key_name, key)
        return key
    
    def generate_iv(self, iv_size=16, iv_name=None):
        """
        生成新的随机IV。
        
        参数:
            iv_size (int): IV大小（字节）
            iv_name (str): 保存IV的可选名称
            
        返回:
            bytes: 生成的IV
        """
        iv = get_random_bytes(iv_size)
        if iv_name:
            self.save_iv(iv_name, iv)
        return iv
    
    def derive_key(self, password, salt=None, key_size=32):
        """
        使用PBKDF2从密码派生密钥。
        
        参数:
            password (str): 用于派生密钥的密码
            salt (bytes): 可选的盐值（如果为None则生成新的）
            key_size (int): 期望的密钥大小（字节）
            
        返回:
            tuple: (派生的密钥字节, 盐值字节)
        """
        if salt is None:
            salt = os.urandom(16)
            
        kdf = PBKDF2HMAC(
            algorithm=hashes.SHA256(),
            length=key_size,
            salt=salt,
            iterations=100000,
        )
        
        key = kdf.derive(password.encode())
        return key, salt
    
    def save_key(self, key_name, key):
        """
        保存带有相关名称的密钥。
        
        参数:
            key_name (str): 用于标识密钥的名称
            key (bytes): 要保存的密钥
        """
        if isinstance(key, bytes):
            key = base64.b64encode(key).decode('utf-8')
        if key_name not in self.keys:
            self.keys[key_name] = {}
        self.keys[key_name]['key'] = key
        self._save_keys()
    
    def save_iv(self, iv_name, iv):
        """
        保存带有相关名称的IV。
        
        参数:
            iv_name (str): 用于标识IV的名称
            iv (bytes): 要保存的IV
        """
        if isinstance(iv, bytes):
            iv = base64.b64encode(iv).decode('utf-8')
        if iv_name not in self.keys:
            self.keys[iv_name] = {}
        self.keys[iv_name]['iv'] = iv
        self._save_keys()
    
    def load_key(self, key_name):
        """
        通过名称加载保存的密钥。
        
        参数:
            key_name (str): 要加载的密钥名称
            
        返回:
            bytes: 加载的密钥
            
        异常:
            KeyError: 如果找不到key_name
        """
        if key_name not in self.keys or 'key' not in self.keys[key_name]:
            raise KeyError(f"找不到密钥 '{key_name}'")
        
        key_str = self.keys[key_name]['key']
        return base64.b64decode(key_str)
    
    def load_iv(self, iv_name):
        """
        通过名称加载保存的IV。
        
        参数:
            iv_name (str): 要加载的IV名称
            
        返回:
            bytes: 加载的IV
            
        异常:
            KeyError: 如果找不到iv_name
        """
        if iv_name not in self.keys or 'iv' not in self.keys[iv_name]:
            raise KeyError(f"找不到IV '{iv_name}'")
        
        iv_str = self.keys[iv_name]['iv']
        return base64.b64decode(iv_str)
    
    def export_key(self, key_name, export_file):
        """
        将密钥导出到文件。
        
        参数:
            key_name (str): 要导出的密钥名称
            export_file (str): 导出密钥的路径
            
        异常:
            KeyError: 如果找不到key_name
        """
        key = self.load_key(key_name)
        with open(export_file, 'wb') as f:
            f.write(key)
    
    def import_key(self, key_name, import_file):
        """
        从文件导入密钥。
        
        参数:
            key_name (str): 保存导入密钥的名称
            import_file (str): 导入密钥的路径
        """
        with open(import_file, 'rb') as f:
            key = f.read()
        self.save_key(key_name, key)
    
    def list_keys(self):
        """
        列出所有保存的密钥名称。
        
        返回:
            list: 所有保存的密钥名称列表
        """
        return [name for name in self.keys.keys() if 'key' in self.keys[name]]
    
    def list_ivs(self):
        """
        列出所有保存的IV名称。
        
        返回:
            list: 所有保存的IV名称列表
        """
        return [name for name in self.keys.keys() if 'iv' in self.keys[name]]
    
    def delete_key(self, key_name):
        """
        删除保存的密钥。
        
        参数:
            key_name (str): 要删除的密钥名称
            
        异常:
            KeyError: 如果找不到key_name
        """
        if key_name in self.keys and 'key' in self.keys[key_name]:
            del self.keys[key_name]['key']
            if not self.keys[key_name]:
                del self.keys[key_name]
            self._save_keys()
        else:
            raise KeyError(f"找不到密钥 '{key_name}'")
    
    def delete_iv(self, iv_name):
        """
        删除保存的IV。
        
        参数:
            iv_name (str): 要删除的IV名称
            
        异常:
            KeyError: 如果找不到iv_name
        """
        if iv_name in self.keys and 'iv' in self.keys[iv_name]:
            del self.keys[iv_name]['iv']
            if not self.keys[iv_name]:
                del self.keys[iv_name]
            self._save_keys()
        else:
            raise KeyError(f"找不到IV '{iv_name}'")